from querier.esquerier import ElasticSearchQuerier
from utils.kmeans import kmeans
from utils import utils
from utils.hc import hc
from utils.hc import hc2


class WechatTitleClusteringQuerier(ElasticSearchQuerier):

    def __init__(self, article_search_query, nlp_service=None):
        super(WechatTitleClusteringQuerier, self).__init__(None, None, None)
        self.article_search_query = article_search_query
        self.nlp_service = nlp_service

    def search(self, args):
        term = args.get('term', '')
        term = term if term else ''
        filters = args.get('filters', {})
        if filters is None:
            filters = {}
        order = args.get('order_by', utils.ORDER_OVERALL)
        from_ = args.get('from', 0)
        if from_ is None:
            from_ = 0
        size_ = args.get('size', 10)
        analyze = args.get('analyze', False)
        n_cluster = args.get('n_cluster')
        filter_keywords = args.get('filter_keywords')
        deduplicate = args.get('deduplicate', False)
        highlight = args.get('highlight', False)
        method = args.get('method')

        if deduplicate is None:
            deduplicate = False

        if not n_cluster:
            n_cluster = 10

        if analyze:
            size_ = 1000

        search_args = {
            'term': term,
            'filters': filters,
            'order_by': order,
            'from': from_,
            'size': size_,
            'deduplicate': deduplicate,
            'highlight': highlight,
            'filter_keywords': filter_keywords
        }

        res = self.article_search_query.search(search_args)
        if not analyze:
            return res
        else:
            titles, reads, urls, summaries = self.get_titles(res)

            try:
                if method == 'hc':
                    out = hc(titles, reads, urls, self.nlp_service, n_cluster)
                elif method == 'kmeans':
                    out = kmeans(titles, reads, urls, n_cluster)
                else:
                    out = hc2(titles, reads, urls, summaries, n_cluster)

            except Exception as e:
                print('%s: %s' % (__file__, str(e)))
                out = []

            return {'title_cluster': out}

    @staticmethod
    def get_titles(article_res):
        titles = []
        reads = []
        urls = []
        summaries = []
        for d in article_res.get('articles', []):
            titles.append(d.get('title', ''))
            reads.append(d.get('read_num'))
            urls.append(d.get('url'))
            summaries.append(d.get('abstract'))
        return titles, reads, urls, summaries

    def _build_query(self, args):
        pass

    def _build_result(self, es_result, param):
        pass
